"""
自定义数据迭代器，即自定义生成器，从原始文件中读取所有的数据，然后按照指定条数，生成每批次的数据
目的：为了后续的ai模型训练课程做铺垫，后续训练模型的时候，不是一次性投喂给大批量的数据，而是分批次来训练的
"""
import math


# 需求：读取jaychou_lyrics.txt文件中的数据，按照n条/批次，生成生成器，并测试
def dataset_loader(batch_size):
    # 读取周杰伦歌词数据文件
    # 打开指定路径的文本文件，以只读模式读取所有行内容
    # 1: 读取文件，获取所有的数据，即：readlines()一次性读取所有行数据，放到列表中
    with open("../data/jaychou_lyrics.txt","r", encoding="utf-8") as src_file:
        data_lines = src_file.readlines()

    # 2: 计算：数据的总条数，即：len(data_lines)
    line_count = len(data_lines)

    # 3: 计算：数据的总批次数，即：line_count // batch_size
    batch_count = math.ceil(line_count/batch_size)

    # 4: 遍历：数据总批次数，获取到具体的每个批次的编号，例如：一共五个批次，那就是：0,1,2,3,4
    for i in range(batch_count):
        # 5: 获取到当前批次的起始索引，例如：0,5,10,15,20
        start_index = i * batch_size

        # 6: 获取到当前批次的结束索引，例如：5,10,15,20,25
        end_index = (i + 1) * batch_size

        # 7: 获取到当前批次的数据，例如：data_lines[0:5]
        batch_data = data_lines[start_index:end_index] # data_lines[0:5]、data_lines[5:10]

        # 8: 返回当前批次的数据
        yield batch_data

if __name__ == '__main__':
    # 测试上述函数，获取指定条数的批次数据。
    my_generator = dataset_loader(5)

    # 获取第一个批次的数据
    # print(next(my_generator))
    # print(next(my_generator))

    for i in my_generator:
        print(i)

